In [1]:
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from scipy import constants

import sys
!{sys.executable} -m pip install plotly

a0 = constants.physical_constants['Bohr radius'][0]
Angstrom = constants.angstrom

Requirement already satisfied: plotly in c:\users\alex1\anaconda3\lib\site-packages (5.9.0)
Requirement already satisfied: tenacity>=6.2.0 in c:\users\alex1\anaconda3\lib\site-packages (from plotly) (8.2.2)
In [2]:
def r_phi_theta(x, y, z):
    r = np.sqrt(x**2 + y**2 + z**2)
    phi = np.arctan2(y, x) + np.pi
    theta = np.arccos(z/r)
    return r, phi, theta
In [3]:
X, Y, Z = np.mgrid[-20*a0:20*a0:48j, -20*a0:20*a0:48j, -20*a0:20*a0:48j]
r_grid, phi_grid, theta_grid = r_phi_theta(X, Y, Z)
In [4]:
def probfunc(r, wavefunc):
    p = r**2 * np.conj(wavefunc) * wavefunc
    return p.real / np.max(p.real)
# The .real syntax picks off the real part of a complex number. Even though the 
# imaginary part of the probability is zero, variable p is still represented in
# the computer as a complex number. Our plotting routines will only work on
# arrays of real numbers.

Question 1: How come these wave functions do not depend on $\phi$ and $\theta$?

In [11]:
# n = 3, l = 2, m = 2
n3l2m2 = lambda r, phi, theta: (r**2/a0**2) * np.exp(-r/(3*a0)) * np.sin(theta) * \
                               np.exp(2j * phi) / \
                               (162 * np.sqrt(np.pi) * a0**1.5)

prob_n3l2m2 = probfunc(r_grid, n3l2m2(r_grid, phi_grid, theta_grid))

fig = go.Figure(data=go.Isosurface(
    x=X.flatten() / Angstrom, # x values of the grid in Angstroms
    y=Y.flatten() / Angstrom, # y values of the grid in Angstroms
    z=Z.flatten() / Angstrom, # z values of the grid in Angstroms
    value=prob_n3l2m2.flatten(), # independent variable
    isomin=0.05, # Minimum normalized probability density to render in an isosurface
    isomax=0.95, # Maximum normalized probability density to render in an isosurface
    opacity=0.4, # Set a low opacity so each surface is partially transparent
    colorscale='Plotly3_r', # Nice-looking color table
    surface_count=8, # number of isosurfaces to plot (2 by default: only min and max)
    colorbar_nticks=8, # colorbar ticks correspond to isosurface values
    caps=dict(x_show=False, y_show=False, z_show=False))
               )

# Change axis lables and make the plot larger than the default
fig.update_layout(scene = dict(
                    xaxis_title='x (Angstroms)',
                    yaxis_title='y (Angstroms)',
                    zaxis_title='z (Angstroms)'),
                    width=700,
                    margin=dict(r=10, b=10, l=10, t=10))

fig.show()
In [ ]: